package com.menghao.rpc.zookeeper;

import lombok.Getter;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.recipes.cache.PathChildrenCache;
import org.apache.curator.framework.recipes.cache.PathChildrenCacheEvent;
import org.apache.curator.framework.recipes.cache.PathChildrenCacheListener;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.zookeeper.CreateMode;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.Assert;

import java.util.List;

/**
 * <p>Zk客户端.</br>
 *
 * @author MarvelCode
 */
public class CuratorClient {

    private final static Logger LOG = LoggerFactory.getLogger(CuratorClient.class);

    private CuratorFramework zkClient;

    @Getter
    private String rootPath;

    public CuratorClient(String zkServerHost, int sessionTimeout, int connectTimeout, String rootPath) {
        if (rootPath.startsWith("/")) {
            rootPath = rootPath.substring(1, rootPath.length());
        }
        zkClient = CuratorFrameworkFactory.builder()
                .connectString(zkServerHost)
                // 会话超时时间，默认60000ms
                .sessionTimeoutMs(sessionTimeout)
                // 连接超时时间
                .connectionTimeoutMs(connectTimeout)
                // 重试策略：每1s重试一次，最大重试次数3
                .retryPolicy(new ExponentialBackoffRetry(1000, 3))
                // 即操作目录
                .namespace(rootPath)
                .build();
        zkClient.start();
        this.rootPath = rootPath;
    }

    public void close() {
        zkClient.close();
    }

    public boolean isExisted(String path) throws Exception {
        return null != zkClient.checkExists().forPath(path);
    }

    public void createEphemeralNode(String path, String value) throws Exception {
        path = checkPath(path);
        if (value != null) {
            zkClient.create().creatingParentsIfNeeded()
                    .withMode(CreateMode.EPHEMERAL)
                    .forPath(path, value.getBytes("UTF-8"));
        } else {
            zkClient.create().creatingParentsIfNeeded()
                    .withMode(CreateMode.EPHEMERAL)
                    .forPath(path, null);
        }
    }

    public List<String> getChildren(String path) throws Exception {
        path = checkPath(path);
        return zkClient.getChildren().forPath(path);
    }

    public void addPathListener(String path, final ChildChangeListener childChangeListener) throws Exception {
        final String correctPath = checkPath(path);
        PathChildrenCache pathChildrenCache = new PathChildrenCache(zkClient, correctPath, true);
        // 永久监听指定节点下一级节点变化
        pathChildrenCache.getListenable().addListener(new PathChildrenCacheListener() {
            @Override
            public void childEvent(CuratorFramework curatorFramework,
                                   PathChildrenCacheEvent pathChildrenCacheEvent) throws Exception {
                if (pathChildrenCacheEvent.getData() != null) {
                    String changePath = pathChildrenCacheEvent.getData().getPath();
                    byte[] changeByte = pathChildrenCacheEvent.getData().getData();
                    String changeData = null;
                    if (changeByte != null) {
                        changeData = new String(changeByte, "UTF-8");
                    }
                    switch (pathChildrenCacheEvent.getType()) {
                        case CHILD_ADDED:
                        case CHILD_UPDATED:
                        case CHILD_REMOVED: {
                            childChangeListener.onChange(changePath, changeData);
                            break;
                        }
                        default:
                            break;
                    }
                }
            }
        });
        pathChildrenCache.start(PathChildrenCache.StartMode.POST_INITIALIZED_EVENT);
    }

    private String checkPath(String path) {
        Assert.hasLength(path, "path must have length");
        if (!path.startsWith("/")) {
            path = "/" + path;
        }
        return path;
    }

}
